#include <iostream>
#include "jcannx.h"
#include "common.h"
#include "matrix_mul.h"

using namespace common;


/*
 * Class:     org_dromara_easyai_extensions_ascend_JCannx
 * Method:    init
 * Signature: ()I
 */
JNIEXPORT jint JNICALL Java_org_dromara_easyai_extensions_ascend_JCannx_init(JNIEnv *env, jclass clz) {
	
	return 0;
}

/*
 * Class:     org_dromara_easyai_extensions_ascend_JCannx
 * Method:    matrixMul
 * Signature: ([F[F[FIII)I
 */
JNIEXPORT jint JNICALL Java_org_dromara_easyai_extensions_ascend_JCannx_matrixMul(JNIEnv *env, jclass clz, jfloatArray xArr, jfloatArray yArr, jfloatArray rArr, jint m, jint k, jint n) {
	if (xArr == nullptr || yArr == nullptr || rArr == nullptr) {
		return 1000;
	}

	// 1、初始化
	int32_t deviceId = 6;
	aclrtStream stream;
	auto ret = Init(deviceId, &stream);
	CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

	// 2、参数准备
	jfloat* x = env->GetFloatArrayElements(xArr, nullptr);
	// printf("Debug: Pointer x address = %p\n", (void*)x);
	// CHECK_RET(x == nullptr, LOG_PRINT("GetFloatArrayElements x failed. ERROR: %d\n", ret); return -1001);

	jfloat* y = env->GetFloatArrayElements(yArr, nullptr);
	// printf("Debug: Pointer y address = %p\n", (void*)y);
	// CHECK_RET(y == nullptr, LOG_PRINT("GetFloatArrayElements y failed. ERROR: %d\n", ret); return -1002);

	jfloat* r = env->GetFloatArrayElements(rArr, nullptr);
	// printf("Debug: Pointer r address = %p\n", (void*)r);
	// CHECK_RET(r == nullptr, LOG_PRINT("GetFloatArrayElements r failed. ERROR: %d\n", ret); return -1003);

	// 3、调用方法计算矩阵相乘
	ret = matrix::Mul(&stream, x, y, r, m, k, n);
	CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("MatMul calculation failed. ERROR: %d\n", ret); return ret);

	// for (int64_t i = 0; i < m * n; i++) {
 //    	printf("result[%ld] is: %f\n", i, r[i]);
 //  	}

	// 4、释放资源
	env->ReleaseFloatArrayElements(xArr, x, 0);
	env->ReleaseFloatArrayElements(yArr, y, 0);
	env->ReleaseFloatArrayElements(rArr, r, 0);

	aclrtDestroyStream(stream);
	aclrtResetDevice(deviceId);
	aclFinalize();
	return 0;
}

/*
 * Class:     org_dromara_easyai_extensions_ascend_JCannx
 * Method:    release
 * Signature: ()I
 */
JNIEXPORT jint JNICALL Java_org_dromara_easyai_extensions_ascend_JCannx_release(JNIEnv *env, jclass clz) {
	
	return 0;
}
